%%%%%%%%%%
% fit_ar.m
%%%%%%%%%%

mdl_plot = 0;

% Data comes in as N x 1 cell array
% need to ensure that each element
% is a row vector, not a col vector
N = length(Data);
for n = 1:N
  [num_rows, num_cols] = size(Data{n});

  if num_cols > num_rows
    Data{n} = Data{n}';
  end
end

% Compute the mean length of the obs vectors
for n = 1:N
  num_samples(n) = length(Data{n});
end
num_samples = mean(num_samples);

% Setup local vars
max_order = 10;
noise_var = zeros(max_order, 1);
pred_error = zeros(max_order, 1);
MDL_term1 = zeros(max_order, 1);
MDL_term2 = zeros(max_order, 1);

% Fit AR model to dataset over a range of model orders
for order=1:max_order
  model{order} = ar(iddata(Data), order, 'burg');
  noise_var(order) = get(model{order}, 'NoiseVariance');
  pred_error(order) = num_samples * noise_var(order);
  MDL_term1(order) = num_samples * log(pred_error(order));
  MDL_term2(order) = order * log(num_samples);
end

% MDL plots for range of orders
if mdl_plot == 1
  plot([1:max_order], MDL_term1, 'b-');
  hold on;
  plot([1:max_order], MDL_term1 + MDL_term2, 'r-');
  legend('AR prediction error', 'AR MDL error');
  hold off;
end

% Find optimal model order
[y,idx] = min(MDL_term1 + MDL_term2);
ar_order = idx(1);
ar_model = model{ar_order};

% Extract AR predictor coeffs
ar_coeffs = get(ar_model, 'ParameterVector');

% Add an extra '1' coefficient at the start
% to get in the correct form for the Matlab
% signal processing filter representation
ar_coeffs = [1;ar_coeffs];

% Compute roots of the AR transfer func
% and also the AR reflection coeffs
ar_roots = roots(ar_coeffs);
ar_rcs = poly2rc(ar_coeffs);

% Now check the AR model stability by
% looking for RCs with a magntiude > 1
for i=1:length(ar_rcs)
  if abs(ar_rcs(i)) > 1.0
    % If so "map" the associated root onto the unit circle
    ar_roots(i) = ar_roots(i)/abs(ar_roots(i));
  end
end

% Convert back to the AR predictor coefficient form
ar_coeffs = poly(ar_roots);

% Remove the first '1' AR coefficient which is simply the
% AR coeff for the signal sample y_{t} we are modelling
ar_coeffs = ar_coeffs(2:end);

% Now, Matlab works with an AR model of the form:
% y_{t} + \sum_{i=1}^{p} a_{i}*y_{t-i} = e_{t}
%
% Therefore to get in the desired form:
% y_{t} = \sum_{i=1}^{p} a_{i}*y_{t-i} + e_{t}
%
% we just need to multiply the AR coeffs by -1

ar_coeffs = -1.0 * ar_coeffs;

% Finally compute AR residual statistics
ar_var = noise_var(ar_order);
ar_mean = 0.0;
ar_precision = inv(ar_var);
ar_norm = sqrt(ar_precision/(2*pi));


save ar_data;
